package com.xu.rag.service.impl.model;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.json.JSONUtil;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.tools.FunctionDefinition;
import com.alibaba.dashscope.tools.ToolCallBase;
import com.alibaba.dashscope.tools.ToolCallFunction;
import com.alibaba.dashscope.tools.ToolFunction;
import com.alibaba.dashscope.utils.Constants;
import com.alibaba.dashscope.utils.JsonUtils;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.github.victools.jsonschema.generator.*;
import com.google.gson.Gson;
import com.xu.rag.common.annotation.ModelPoint;
import com.xu.rag.common.annotation.ModelType;
import com.xu.rag.config.LLmProperties;
import com.xu.rag.functionCall.GetNamesTool;
import com.xu.rag.functionCall.GetTimeTool;
import com.xu.rag.functionCall.GetWhetherTool;
import com.xu.rag.pojo.result.ChunkResult;
import com.xu.rag.pojo.result.EmbeddingResult;
import com.xu.rag.pojo.result.TongyiResponse;
import com.xu.rag.pojo.result.TongyiResult;
import com.xu.rag.service.LlmService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * @Author: xdl
 * @CreateTime: 2025-01-13 23:22
 * @Description: 通义千问
 */
@Slf4j
@Service
@ModelType(modelCode = "TONG_YI")
public class TongyiServiceImpl implements LlmService {

    private static final String GET_CURRENT_WHETHER = "get_current_whether";
    private static final String GET_CURRENT_TIME = "get_current_time";
    private static final String GET_CURRENT_NAMES = "get_current_names";
    @Resource
    private LLmProperties lLmProperties;

    @Override
    public String functionCall(String message) throws Exception {
        Constants.apiKey = lLmProperties.getTongyiApiKey();
        SchemaGeneratorConfigBuilder configBuilder = new SchemaGeneratorConfigBuilder(SchemaVersion.DRAFT_2020_12, OptionPreset.PLAIN_JSON);
        SchemaGeneratorConfig config = configBuilder.with(Option.EXTRA_OPEN_API_FORMAT_VALUES)
                .without(Option.FLATTENED_ENUMS_FROM_TOSTRING)
                .build();
        SchemaGenerator generator = new SchemaGenerator(config);
        ObjectNode jsonSchema_whether = generator.generateSchema(GetWhetherTool.class);
        ObjectNode jsonSchema_time = generator.generateSchema(GetTimeTool.class);
        ObjectNode jsonSchema_names = generator.generateSchema(GetNamesTool.class);

        FunctionDefinition fd_whether = FunctionDefinition.builder()
                .name(GET_CURRENT_WHETHER)
                .description("获取指定地区指定时间的天气")
                .parameters(JsonUtils.parseString(jsonSchema_whether.toString()).getAsJsonObject())
                .build();
        FunctionDefinition fd_time = FunctionDefinition.builder()
                .name(GET_CURRENT_TIME)
                .description("获取当前时刻时间")
                .parameters(JsonUtils.parseString(jsonSchema_time.toString()).getAsJsonObject())
                .build();
        FunctionDefinition fd_names = FunctionDefinition.builder()
                .name(GET_CURRENT_NAMES)
                .description("获取好友名称")
                .parameters(JsonUtils.parseString(jsonSchema_names.toString()).getAsJsonObject())
                .build();

        com.alibaba.dashscope.common.Message systemMsg = com.alibaba.dashscope.common.Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你是一个乐于助人的AI助手。当被问到问题时，尽可能使用工具。")
                .build();
        com.alibaba.dashscope.common.Message userMsg = com.alibaba.dashscope.common.Message.builder()
                .role(Role.USER.getValue())
                .content(message)
                .build();

        List<com.alibaba.dashscope.common.Message> messages = new ArrayList<>();
        messages.add(systemMsg);
        messages.add(userMsg);

        GenerationParam param = GenerationParam.builder()
                .model("qwen-turbo")
                .messages(messages)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .tools(Arrays.asList(
                        ToolFunction.builder().function(fd_whether).build(),
                        ToolFunction.builder().function(fd_time).build(),
                        ToolFunction.builder().function(fd_names).build()))
                .build();

        Generation gen = new Generation();
        GenerationResult result = gen.call(param);

        for (GenerationOutput.Choice choice : result.getOutput().getChoices()) {
            messages.add(choice.getMessage());
            List<ToolCallBase> toolCalls = result.getOutput().getChoices().get(0).getMessage().getToolCalls();
            if (toolCalls != null) {
                for (ToolCallBase toolCall : toolCalls) {
                    if ("function".equals(toolCall.getType())) {
                        String functionName = ((ToolCallFunction) toolCall).getFunction().getName();
                        String functionArgument = ((ToolCallFunction) toolCall).getFunction().getArguments();
                        if (GET_CURRENT_WHETHER.equals(functionName)) {
                            GetWhetherTool getWhetherFunction = JsonUtils.fromJson(functionArgument, GetWhetherTool.class);
                            String whether = getWhetherFunction.call();
                            com.alibaba.dashscope.common.Message toolResultMessage = com.alibaba.dashscope.common.Message.builder()
                                    .role("tool")
                                    .content(String.valueOf(whether))
                                    .toolCallId(toolCall.getId())
                                    .build();
                            messages.add(toolResultMessage);
                        } else if (GET_CURRENT_TIME.equals(functionName)) {
                            GetTimeTool getTimeFunction = JsonUtils.fromJson(functionArgument, GetTimeTool.class);
                            String time = getTimeFunction.call();
                            com.alibaba.dashscope.common.Message toolResultMessage = com.alibaba.dashscope.common.Message.builder()
                                    .role("tool")
                                    .content(String.valueOf(time))
                                    .toolCallId(toolCall.getId())
                                    .build();
                            messages.add(toolResultMessage);
                        } else if (GET_CURRENT_NAMES.equals(functionName)) {
                            GetNamesTool getNamesFunction = JsonUtils.fromJson(functionArgument, GetNamesTool.class);
                            String names = getNamesFunction.call();
                            com.alibaba.dashscope.common.Message toolResultMessage = com.alibaba.dashscope.common.Message.builder()
                                    .role("tool")
                                    .content(String.valueOf(names))
                                    .toolCallId(toolCall.getId())
                                    .build();
                            messages.add(toolResultMessage);
                        }
                    }
                }
            } else {
                // 无需调用工具
                return "最终答案：" + result.getOutput().getChoices().get(0).getMessage().getContent();
            }
        }
        param.setMessages(messages);
        GenerationResult call = gen.call(param);
        return "最终答案：" + call.getOutput().getChoices().get(0).getMessage().getContent();
    }

    @Override
    public String chat(String prompt) {
        try {
            RequestBody requestBody = new RequestBody(
                    lLmProperties.getTongyiChatModel(),
                    new Message[]{
                            new Message("system", "You are a helpful assistant."),
                            new Message("user", prompt)
                    }
            );
            Gson gson = new Gson();
            String jsonInputString = gson.toJson(requestBody);
            URL url = new URL(lLmProperties.getTongyiChatUrl());
            HttpURLConnection httpURLConnection = (HttpURLConnection) url.openConnection();
            httpURLConnection.setRequestMethod("POST");
            httpURLConnection.setRequestProperty("Content-Type", "application/json; utf-8");
            httpURLConnection.setRequestProperty("Accept", "application/json");
            String auth = "Bearer " + lLmProperties.getTongyiApiKey();
            httpURLConnection.setRequestProperty("Authorization", auth);
            httpURLConnection.setDoOutput(true);
            try (OutputStream os = httpURLConnection.getOutputStream()) {
                byte[] input = jsonInputString.getBytes(StandardCharsets.UTF_8);
                os.write(input, 0, input.length);
            }
            int responseCode = httpURLConnection.getResponseCode();
            try (BufferedReader br = new BufferedReader(new InputStreamReader(httpURLConnection.getInputStream(), StandardCharsets.UTF_8))) {
                StringBuilder response = new StringBuilder();
                String responseLine;
                while ((responseLine = br.readLine()) != null) {
                    response.append(responseLine.trim());
                }
                String string = response.toString();
                TongyiResponse result = JSONUtil.toBean(string, TongyiResponse.class);
                String content = result.getChoices().get(0).getMessage().getContent();
                return content;
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
        return "";
    }

    @Override
//    @ModelPoint
    public double[] sentence(String sentence) {
        ChunkResult chunkResult = new ChunkResult()
                .setContent(sentence)
                .setChunkId(RandomUtil.randomInt());
        return embedding(chunkResult).getEmbedding();
    }

    @Override
    public List<EmbeddingResult> embedding(List<ChunkResult> chunkResults) {
        if (CollectionUtil.isEmpty(chunkResults)) {
            return new ArrayList<>();
        }
        List<EmbeddingResult> embeddingResults = new ArrayList<>();
        for (ChunkResult chunkResult : chunkResults) {
            embeddingResults.add(embedding(chunkResult));
        }
        return embeddingResults;
    }

    @Override
    public EmbeddingResult embedding(ChunkResult chunkResult) {
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", lLmProperties.getTongyiEmbeddingModel());
        requestBody.put("input", chunkResult.getContent());
        requestBody.put("dimensions", 1024);
        requestBody.put("encoding_format", "float");
        try {
            String requestBodyString = JSONUtil.toJsonStr(requestBody);
            String response = HttpRequest.post(lLmProperties.getTongyiEmbeddingUrl())
                    .header("Content-Type", "application/json")
                    .header("Authorization", "Bearer" + lLmProperties.getTongyiApiKey())
                    .body(requestBodyString)
                    .execute()
                    .body();
            TongyiResult tongyiResult = JSONUtil.toBean(response, TongyiResult.class);
            List<TongyiResult.EmbeddingData> data = tongyiResult.getData();
            return new EmbeddingResult()
                    .setEmbedding(data.get(0).getEmbedding())
                    .setPrompt(chunkResult.getContent())
                    .setRequestId(String.valueOf(chunkResult.getChunkId()));
        } catch (Exception e) {
            log.error("Embedding error: {}", e.getMessage());
        }
        return null;
    }

    static class Message {
        String role;
        String content;

        public Message(String role, String content) {
            this.role = role;
            this.content = content;
        }
    }

    static class RequestBody {
        String model;
        Message[] messages;

        public RequestBody(String model, Message[] messages) {
            this.model = model;
            this.messages = messages;
        }
    }

}
